import matplotlib.pyplot as plt
import os
import torch
from torch.utils.data import DataLoader
from torch import nn
import math
import csv


class SavingMethod:
    def __init__(self,model_name,path = './img'):
        self.createFolder(path)
        self.save_path = path + '/' + str(len(os.listdir(path)) )
        os.mkdir(self.save_path)
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        #self.device = 'cpu'
        self.Datapath = self.save_path + '/' + 'data' + str(len(os.listdir(self.save_path))) + '.csv'
        self.lossPath = self.save_path + '/' + 'loss' + str(len(os.listdir(self.save_path))) + '.csv'
        self.model_name = model_name
        self.saveCode()
        self.RMSE = []
        self.MAE = []
        self.SCORE = []
        self.LOSS  = []
        self.models = {}


    def createFolder(self,path):
        if not os.path.exists(path):
            os.mkdir(path)
    ##保存所有的loss 包括第几个loader，第几个epoch 第几个训练，loss
    def savetrain_loss(self,loss:list):
        self.LOSS.append(loss)
    ##这个函数用于测试集输入和输出的图像并且将其绘制出来
    def savefig(self, model:nn.Module, test_set:DataLoader, epoch:int):
        savePath = self.save_path + '/epoch' + str(epoch) + '.png'
        sum_RMSE_loss = 0
        sum_MAE_loss = 0
        sum_RUL_score = 0
        m = 0
        pltx = []  # 实际值
        plty = []  # 预测值
        total_loss = 0
        temp_num = math.log(0.5)

        for i,(x,y,label) in enumerate(test_set):
            x = x.to(self.device)
            label = label.to(self.device)
            label_hat = model(x)  #此时label_hat是一个二维的数组（batch_size,1)
            label_hat = label_hat.tolist()
            label = label.tolist()
            t = 0
            m += len(label_hat)
            for x in range(len(label_hat)):
                RULr = label[x][0]
                RULp = label_hat[x][0]
                pltx.append(RULr)
                plty.append(RULp)

                #  计算分数
                sum_RMSE_loss += (RULr- RULp) ** 2
                sum_MAE_loss += abs(RULr- RULp)
                E = (RULr- RULp ) * 100 / RULr
                if E <= 0:
                    sum_RUL_score += math.exp(-temp_num*(E / 5))
                else:
                    sum_RUL_score += math.exp(temp_num*(E / 20))
                t += abs(RULr- RULp)
            total_loss += t / len(label_hat)
        #保存数据
        sum_RMSE_loss = math.sqrt(sum_RMSE_loss / m)
        self.RMSE.append(sum_RMSE_loss)
        sum_MAE_loss = sum_MAE_loss / m
        self.MAE.append(sum_MAE_loss)
        sum_RUL_score /= m
        self.SCORE.append(sum_RUL_score)
        #self.savetrain_loss([epoch, total_loss])
        self.saveFile()
        print('epoch:{}, RMSE:{}, MAE:{}, RUL_SCORE:{}'.format(epoch, sum_RMSE_loss, sum_MAE_loss, sum_RUL_score))
        #保存图片
        plt.clf()
        plt.plot(range(len(pltx)), pltx, c='blue',label = 'RUL truth')
        plt.scatter(range(len(pltx)), plty, c='red', s=0.2,label = 'RUL predict')
        plt.legend()
        plt.savefig(savePath)
        self.saveModel(epoch,sum_RMSE_loss,model)

    def saveModel(self,epoch,rmse,model):
        if len(self.models) < 10:
            self.models[str(epoch)] = rmse
        else:
            dic = sorted(self.models.items(), key=lambda x: x[1])
            if dic[-1][1] > rmse:
                self.models.pop(dic[-1][0])
                os.remove('./trained_model/' + str(dic[-1][0]) + '.pt')
                self.models[str(epoch)] = rmse
            else:
                print('epoch:{}, rmse:{} model has been dropped'.format(epoch, rmse))
                return
        print('epoch:{}, rmse:{} model has been appended'.format(epoch,rmse))
        torch.save(model, './trained_model/' + str(epoch) + '.pt')

    def saveFile(self):
        with open(self.Datapath, 'w', newline='')as f:
            writer1 = csv.writer(f)
            writer1.writerow(['epoch', 'RMSE', 'MAE', 'RUL_SCORE'])
            for i in range(len(self.MAE)):
                writer1.writerow([i,self.RMSE[i],self.MAE[i],self.SCORE[i]])
        with open(self.lossPath, 'w', newline='')as f:
            writer2 = csv.writer(f)
            writer2.writerow(['epoch', 'condition', 'Bearing', 'time', 'all_time', 'loss'])
            writer2.writerows(self.LOSS)

    def saveCode(self):
        import shutil
        #shutil.copy('./Save.py',self.save_path)
        #shutil.copy('./loader.py', self.save_path)
        for name in self.model_name:
            shutil.copy(name, self.save_path)


if __name__ == '__main__':
    main